{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "分\n壕\n"
    }
   ],
   "source": [
    "import time\n",
    "import math\n",
    "import numpy as np \n",
    "import torch\n",
    "import torch.nn as nn \n",
    "import torch.nn.functional as F \n",
    "import d2lzh as d2l\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "corpus_indices,char_to_idx,idx_to_char,vocab_size = d2l.load_data_jay_lyrics()\n",
    "\n",
    "print(idx_to_char[char_to_idx[\"分\"]])\n",
    "print(idx_to_char[25])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## one-hot向量\n",
    "为了将**词**表示成向量输入到网络中，一个简单的办法市使用one-hot向量。　假设词典中不同字符的个数为$N$（即词典的大小），每个字符已经和一个从0到$N-1$的连续整数索引一一对应。如果一个字符的索引市整数$i$，那么可以创建一个全０的长度为$N$的向量，并将其位置为$i$的元素设置为１．　该向量就是原字符的one-hot向量."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n        [0., 0., 1.,  ..., 0., 0., 0.]])"
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "source": [
    "def one_hot(x,n_class,dtype=torch.float32):\n",
    "    x = x.long()\n",
    "    res = torch.zeros(x.shape[0],n_class,dtype=dtype,device=x.device)\n",
    "    res.scatter_(1,x.view(-1,1),1)\n",
    "    return res\n",
    "\n",
    "x = torch.tensor([0,2])\n",
    "one_hot(x,vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "5 torch.Size([2, 1447])\n"
    }
   ],
   "source": [
    "def to_onehot(x,n_class):\n",
    "    # x shape:(batch,seq_len),\n",
    "    # output: seq_len elements of (batch,n_class)\n",
    "\n",
    "    return [one_hot(x[:,i],n_class) for i in range(x.shape[1])]\n",
    "\n",
    "x = torch.arange(10).view(2,5)\n",
    "inputs = to_onehot(x,vocab_size)\n",
    "print(len(inputs),inputs[0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "use  cuda\n"
    }
   ],
   "source": [
    "num_inputs,num_hiddens,num_outputs = vocab_size,256,vocab_size\n",
    "print(\"use \",device)\n",
    "\n",
    "def get_params():\n",
    "\n",
    "    def _one(shape):\n",
    "        ts = torch.tensor(np.random.normal(0,0.01,size=shape),device=device,dtype=torch.float32)\n",
    "        return torch.nn.Parameter(ts,requires_grad=True)\n",
    "    \n",
    "    w_xh = _one((num_inputs,num_hiddens))\n",
    "    w_hh = _one((num_hiddens,num_hiddens))\n",
    "    b_h = torch.nn.Parameter(torch.zeros(num_hiddens,device=device,requires_grad=True))\n",
    "\n",
    "    #　输出层\n",
    "    w_hq = _one((num_hiddens,num_outputs))\n",
    "    b_q = torch.nn.Parameter(torch.zeros(num_outputs,device=device,requires_grad=True))\n",
    "\n",
    "    return nn.ParameterList([w_xh,w_hh,b_h,w_hq,b_q])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "5 torch.Size([2, 1447]) torch.Size([2, 256])\n"
    }
   ],
   "source": [
    "def init_rnn_state(batch_size,num_hiddens,device):\n",
    "    return (torch.zeros((batch_size,num_hiddens),device=device),)\n",
    "\n",
    "def rnn(inputs,state,params):\n",
    "    w_xh,w_hh,b_h,w_hq,b_q = params \n",
    "    H, = state\n",
    "    outputs = []\n",
    "    for x in inputs:\n",
    "        H = torch.tanh(torch.matmul(x,w_xh) + torch.matmul(H,w_hh) + b_h)\n",
    "        y = torch.matmul(H,w_hq) + b_q\n",
    "        outputs.append(y)\n",
    "    return outputs,(H,)\n",
    "\n",
    "state = init_rnn_state(x.shape[0],num_hiddens,device)\n",
    "inputs = to_onehot(x.to(device),vocab_size)\n",
    "params = get_params()\n",
    "outputs,state_new = rnn(inputs,state,params)\n",
    "print(len(outputs),outputs[0].shape,state_new[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 预测函数\n",
    "基于前缀`prefix`（含有数个字符的字符串）来预测接下来的`num_chars`个字符。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_rnn(prefix,num_chars,rnn,params,init_rnn_state,num_hiddens,vocab_size,device,idx_to_char,char_to_idx):\n",
    "\n",
    "    state = init_rnn_state(1,num_hiddens,device)\n",
    "    output = [char_to_idx[prefix[0]]]\n",
    "    for t in range(num_chars + len(prefix) - 1):\n",
    "        # 将上一个时间步的输出作为下一个时间步的输入\n",
    "        x = to_onehot(torch.tensor([[output[-1]]],device=device),vocab_size)\n",
    "        (y,state) = rnn(x,state,params) #计算和更新隐藏状态\n",
    "        #下一个时间步的输入是prefix里的字符或者当前的最佳预测字符\n",
    "        if t < len(prefix)-1:\n",
    "            output.append(char_to_idx[prefix[t+1]])\n",
    "        else:\n",
    "            output.append(int(y[0].argmax(dim=1).item())) \n",
    "    return \" \".join([idx_to_char[i] for i in output])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'分 开 练 散 载 至 腾 仪 箱 记 熬 屉'"
     },
     "metadata": {},
     "execution_count": 7
    }
   ],
   "source": [
    "predict_rnn(\"分开\",10,rnn,params,init_rnn_state,num_hiddens,vocab_size,device,idx_to_char,char_to_idx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grad_clipping(params,theta,device):\n",
    "    norm = torch.tensor([0.0],device=device)\n",
    "    for param in params:\n",
    "        norm += (param.grad.data ** 2).sum()\n",
    "    norm = norm.sqrt().item()\n",
    "    if norm > theta:\n",
    "        for param in params:\n",
    "            param.grad.data *= (theta/norm)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_predict_rnn(rnn,get_params,init_rnn_state,num_hiddens,vocab_size,device,corpus_indices,idx_to_char,char_to_idx,is_random_iter,num_epochs,num_steps,lr,clipping_theta,batch_size,pred_period,pred_len,prefixes):\n",
    "    if is_random_iter:\n",
    "        data_iter_fn = d2l.data_iter_random\n",
    "    else:\n",
    "        data_iter_fn = d2l.data_iter_consecutive\n",
    "\n",
    "    params = get_params()\n",
    "    loss = nn.CrossEntropyLoss()\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        if not is_random_iter: # 使用相邻采样，在epoch开始时初始化隐藏层\n",
    "            state = init_rnn_state(batch_size,num_hiddens,device)\n",
    "        l_sum,n,start = 0.0,0,time.time()\n",
    "        data_iter = data_iter_fn(corpus_indices,batch_size,num_steps,device)\n",
    "        for x,y in data_iter:\n",
    "            #　随机采样，在每个小批量更新钱初始化隐藏状态\n",
    "            if is_random_iter:\n",
    "                state = init_rnn_state(batch_size,num_hiddens,device) \n",
    "            else:\n",
    "                for s in state:\n",
    "                    s.detach_()\n",
    "            inputs = to_onehot(x,vocab_size)\n",
    "            (outputs,state) = rnn(inputs,state,params)\n",
    "            outputs = torch.cat(outputs,dim=0)\n",
    "            y = torch.transpose(x,0,1).contiguous().view(-1)\n",
    "            l = loss(outputs,y.long())\n",
    "\n",
    "            if params[0].grad is not None:\n",
    "                for param in params:\n",
    "                    param.grad.data.zero_()\n",
    "            l.backward()\n",
    "            grad_clipping(params,clipping_theta,device)\n",
    "            d2l.sgd(params,lr,1)\n",
    "\n",
    "            l_sum += l.item() * y.shape[0]\n",
    "            n += y.shape[0]\n",
    "\n",
    "        if (epoch + 1) % pred_period == 0:\n",
    "            print(\"epoch %d,perplexity %f,time %.2f sec\" % (epoch + 1,math.exp(l_sum / n),time.time() - start))\n",
    "\n",
    "            for prefix in prefixes:\n",
    "                print(\" -\",predict_rnn(prefix,pred_len,rnn,params,init_rnn_state,num_hiddens,vocab_size,device,idx_to_char,char_to_idx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "epoch 50,perplexity 3.565956,time 0.23 sec\n - 分 开 开 始 口 边 拜 淡 油 啦 愿 东\nepoch 100,perplexity 1.346147,time 0.24 sec\n - 分 开 开 始 口 边 弥 淡 淡 … 产 鸣\nepoch 150,perplexity 1.091614,time 0.23 sec\n - 分 开 开 始 口 语 比 方 方 平 产 得\nepoch 200,perplexity 1.037620,time 0.23 sec\n - 分 开 开 始 口 边 天 天 天 天 天 天\nepoch 250,perplexity 1.016120,time 0.23 sec\n - 分 开 开 始 口 口 比 方 方 安 和 和\n"
    }
   ],
   "source": [
    "num_epochs, num_steps, batch_size, lr, clipping_theta = 250, 35, 32, 1e2, 1e-2\n",
    "pred_period, pred_len, prefixes = 50, 10, ['分开']\n",
    "\n",
    "train_and_predict_rnn(rnn, get_params, init_rnn_state, num_hiddens,\n",
    "                      vocab_size, device, corpus_indices, idx_to_char,\n",
    "                      char_to_idx, True, num_epochs, num_steps, lr,\n",
    "                      clipping_theta, batch_size, pred_period, pred_len,\n",
    "                      prefixes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python37764bitpytorchnotebookconda6e7a8693df0d4d92aca09d521275d23a",
   "display_name": "Python 3.7.7 64-bit ('pytorch_notebook': conda)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}